# -*- coding: utf-8 -*-
"""
Created on Wed Oct 27 17:14:27 2021

@author: singularity

https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.ReduceLROnPlateau.html?highlight=reducelronplateau

"""

import torch
import torch.optim as optim


import numpy as np
import PIL
import pandas

import matplotlib.pyplot as plt

import dataset_VID
import yoloJ
import yoloLoss
import glob

import torch.utils.tensorboard as tensorboard

def saveMyYOLO(model,optimizer,savePath:str,epoch:int):
    state = {'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': epoch}
    torch.save(state, savePath)
    print("保存模型到："+savePath)

if __name__ == "__main__":
    BATCH_SIZE_N = 4
    EPOCH = 100
    
    AnchorSzList = [[(50,44), (64,82), (82,48)],
               [(64,114), (108,66), (104,102)]]
    StrideSzList = [8,4]
    
    LOG_PER_N_BATCH = 20
    
    
    '''数据记录'''
    # -- 文件管理
    logPath = './log/DIV-YOLOJ-'
    
    logPath += str(BATCH_SIZE_N)+'BSz'
    
    logDir = glob.glob(logPath+'*')
    
    if len(logDir):
        print("WARN：当前目录下已经存在相关记录！")
        print(logDir)
        cmd = ''
        while cmd == '' or cmd != 'Y' or cmd != 'N':
            cmd = input("是否继续：（Y/N）")
            if cmd =='Y':
                if input("覆盖？（Y/[N]）：") != 'Y':
                    print("NOTE:覆盖原有记录")
                    logPath += '-'
                    logPath += str(len(logDir)+1)
                break
            elif cmd =='N':
                raise Exception("日志已存在且不许覆盖")
            else:
                print("输入命令错误")
        
    print("创建TensorBoard日志在："+logPath)
    
    # -- 创建/打开日志
    writer = tensorboard.SummaryWriter(logPath)
    
    logStep = 1 # 日志step坐标
    preEpochStep = 1
    
    
    '''加载数据'''
    print("开始加载数据")
    trainset= dataset_VID.VID(True)
    trainset.convert2trainData(2, AnchorSzList, StrideSzList)
    trainset.convert2YOLOJ()
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=BATCH_SIZE_N,
                                              shuffle=True, num_workers=2)
    
    testset = dataset_VID.VID(False)
    testset.convert2trainData(2, AnchorSzList, StrideSzList)
    testset.convert2YOLOJ()
    testloader = torch.utils.data.DataLoader(testset, batch_size=BATCH_SIZE_N,
                                              shuffle=True, num_workers=2)
    
    '''
    @gpu
    '''
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(f"计算硬件:{device}")
    
    torch.cuda.empty_cache() # 强行清空显存

    '''建立网络'''
    print("开始建立网络")
    net= yoloJ.YOLOJ()
    net.to(device)
    print("建立网络成功")
    
# =============================================================================
#     print(net)
# =============================================================================
    
    lossF = yoloLoss.YOLOJ_Loss(5, 3)
# =============================================================================
#     optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.8)
# =============================================================================
    
    optimizer = optim.Adam(net.parameters(),
                           lr=0.001,
                           betas=(0.9, 0.999))
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min',
                                                     factor=0.95,    # 衰减率
                                                     patience=10,   # 容忍设为90次batch
                                                     threshold=0.0001,
                                                     min_lr = 1e-6)
    
    
    '''加载预先训练过的模型'''
    if input("是否需要加载已存在的模型：（Y/N）") == 'Y':
        mdPath = input("请输入模型（相对）路径：")
        if len(glob.glob(mdPath)):
            print("----- 开始加载模型 -----")
            checkpoint = torch.load(mdPath)
            net.load_state_dict(checkpoint['model'])
            optimizer.load_state_dict(checkpoint['optimizer'])
            preEpochStep = checkpoint(['epoch'])
    
    
    '''
    @模型训练
    '''
    print("开始训练模型")
    for epoch in range(preEpochStep, EPOCH +preEpochStep):
        print(f"-- 第{epoch}次Epoch --")
        running_loss = 0.0
        
        # 模型进入训练模式
        net.train()
        for i, data in enumerate(trainloader, 0):
            # 数据加载到GPU
            inputs, labels = data[0],data[1]
            inputs = inputs.to(device)
            labels = (labels[0].to(device),labels[1].to(device))
            
            '''训练模型'''
            optimizer.zero_grad()
            # 前向传播
            outputs = net(inputs)
            # (([8, 30, 16, 16]), [8, 30, 32, 32])
            
            #labels: ([8, 33, 16, 16], [8, 33, 32, 32])
            # 损失计算
            loss = lossF(outputs, labels)
            # 反向传播
            loss.backward()
            # 参数更新
            optimizer.step()

            running_loss += loss.item() # 训练损失累积
            
        '''每LOG_PER_N_BATCH次，进行数据记录'''
# =============================================================================
#         if i % LOG_PER_N_BATCH == LOG_PER_N_BATCH-1:
# =============================================================================
        if True:
# =============================================================================
#             running_loss /= LOG_PER_N_BATCH # 训练损失平均
# =============================================================================
            running_loss /= (i+1)            

            '''测试集损失计算'''
# =============================================================================
#             # 清空原有显存占用，没办法，实在太穷了
#             inputs = inputs.cpu()
#             labels = (labels[0].cpu(),labels[1].cpu())
#             outputs = (outputs[0].cpu(),outputs[1].cpu())
#             del outputs
#             loss = loss.cpu()
#             del loss
#             optimizer.zero_grad()
#             
#             torch.cuda.empty_cache()
# =============================================================================
            
            # 模型进入评估模式
            
            net.eval()
            valLoss =0
            with torch.no_grad():
                for iV ,valD in enumerate(testloader):
                    # 数据加载到GPU
                    inputs, labels = valD[0],valD[1]
                    inputs = inputs.to(device)
                    labels = (labels[0].to(device),labels[1].to(device))
                    # 网络前向传播推导
                    outputs = net(inputs)
                    # 损失计算并累加
                    valLoss += lossF(outputs, labels)
                
            valLoss /= (iV+1)  # 损失平均
            
            # 学习率监督调整
            scheduler.step(valLoss)
            
            # 记录数据
            
            print('[%d, %5d] TrainLoss: %.3f , ValLoss: %.3f' %
                  (epoch, i + 1, running_loss,valLoss))
            
            writer.add_scalar("YOLOJ损失/train", running_loss,logStep)
            writer.add_scalar("YOLOJ损失/val", valLoss, logStep)
            writer.add_scalar("学习率", optimizer.state_dict()['param_groups'][0]['lr'],logStep)
            
            logStep +=1
            
            running_loss = 0.0
                